from matplotlib.dates import DateFormatter, \
    WeekdayLocator, DayLocator, MONDAY, date2num
from datetime import datetime
from mpl_finance import candlestick_ohlc
from matplotlib import pyplot as plt


def candlePlot(seriasData, title='a'):
    '''
    seriasData: index  时间序列  Open  High Low Close
    title： 绘图标题
    # df.rename( columns={'times':'Date','openp':'Open', 'highp':'High','lowp':'Low',"nowv":'Close','curvol':"Volume"}, inplace=True)

    '''
    Date = [date2num(date) for date in seriasData.index]
    seriasData.loc[:, 'Date'] = Date

    # DataFrame -> list
    listData = []
    for i in range(len(seriasData)):
        a = [seriasData.Date[i],
             seriasData.Open[i], seriasData.High[i],
             seriasData.Low[i], seriasData.Close[i]]
        listData.append(a)

    # 设定绘图相关参数
    ax = plt.subplot()

    mondays = WeekdayLocator(MONDAY)

    # 15 - Mar - 09
    weekFormatter = DateFormatter("%y %b %d")
    ax.xaxis.set_major_locator(mondays)
    ax.xaxis.set_minor_locator(DayLocator())
    ax.xaxis.set_major_formatter(weekFormatter)

    candlestick_ohlc(ax, listData, width=0.7,
                     colorup='r', colordown='g')

    ax.set_title(title)

    plt.setp(plt.gca().get_xticklabels(),
             rotation=50, horizontalalignment='center')

    return(plt.show())


def PlotKline(seriasData, title='a'):
    
    '''
    seriasData: index  时间序列  open  high low close
    title： 绘图标题
    '''
    Date = [date2num(date) for date in seriasData.index]
    seriasData.loc[:, 'Date'] = Date

    # DataFrame -> list
    listData = []
    for i in range(len(seriasData)):
        a = [seriasData.Date[i],
             seriasData.open[i], seriasData.high[i],
             seriasData.low[i], seriasData.close[i]]
        listData.append(a)

    # 设定绘图相关参数
    ax = plt.subplot()

    mondays = WeekdayLocator(MONDAY)

    # 15 - Mar - 09
    weekFormatter = DateFormatter("%y %b %d")
    ax.xaxis.set_major_locator(mondays)
    ax.xaxis.set_minor_locator(DayLocator())
    ax.xaxis.set_major_formatter(weekFormatter)

    candlestick_ohlc(ax, listData, width=0.7,
                     colorup='r', colordown='g')

    ax.set_title(title)

    plt.setp(plt.gca().get_xticklabels(),
             rotation=50, horizontalalignment='center')
    plt.figure(figsize=(16,8))
    
    return(plt.show())



def candleVolume(df, candletitle='a', bartitle=''):
    
    '''
    seriasData: index  时间序列  Open  High Low Close
    title： 绘图标题
    '''
    from matplotlib import pyplot as plt
    from matplotlib.dates import DateFormatter, WeekdayLocator, DayLocator, MONDAY, date2num
    from mpl_finance import candlestick_ohlc
    import numpy as np
    import warnings

    warnings.filterwarnings('ignore')
    plt.rcParams['font.sans-serif'] = ['SimHei']
    plt.rcParams['axes.unicode_minus'] = False
    plt.figure(figsize=(16, 8))

    seriesData = df.copy()

    Date = [date2num(date) for date in seriesData.index]
    seriesData.index = list(range(len(Date)))
    seriesData['Date'] = Date

    listData = zip(seriesData.Date, seriesData.Open,
                   seriesData.High, seriesData.Low, seriesData.Close)
    ax1 = plt.subplot(211)
    ax2 = plt.subplot(212)
    for ax in ax1, ax2:
        mondays = WeekdayLocator(MONDAY)
        weekFormatter = DateFormatter('%Y%m%d')
        ax.xaxis.set_major_locator(mondays)
        ax.xaxis.set_minor_locator(DayLocator())
        ax.xaxis.set_major_formatter(weekFormatter)
        ax.grid(True)

    ax1.set_ylim(seriesData.Low.min()-1, seriesData.High.max() + 1)
    ax1.set_ylabel('蜡烛图及收盘价线')
    candlestick_ohlc(ax1, listData, width=0.7, colorup='r', colordown='g')
    plt.setp(plt.gca().get_xticklabels(),
             rotation=20, horizontalalignment='center')
    ax1.autoscale_view()
    ax1.set_title(candletitle)

    ax1.plot(seriesData.Date, seriesData.Close, color='black', label='收盘价')
    ax1.legend(loc='best')

    ax2.set_ylabel('成交量')
    ax2.set_ylim(0, seriesData.Volume.max()*2)
    ax2.bar(np.array(Date)[np.array(seriesData.Close >= seriesData.Open)],
            height=seriesData.loc[:, 'Volume'][np.array(
                seriesData.Close >= seriesData.Open)],
            color='r', align='center')
    ax2.bar(np.array(Date)[np.array(seriesData.Close < seriesData.Open)],
            height=seriesData.loc[:, 'Volume'][np.array(
                seriesData.Close < seriesData.Open)],
            color='g', align='center')
    ax2.set_title(bartitle)
    return(plt.show())
